# -*- coding: utf-8 -*-
"""
"""
import torch
import torch.nn as nn
import torch.functional as f
import torchvision
import torchsummary as summary
import torchvision.transforms as transforms
import os
import sys
os.path.join(os.path.dirname(__file__), '../')
sys.path.append(os.path.join(os.path.dirname(__file__), '../'))
from public import newsCNN

'''设置transform'''

# mean_nums = [0.485, 0.456, 0.406]
# std_nums = [0.229, 0.224, 0.225]
mean_nums = [0.5, 0.5, 0.5]
std_nums = [0.5, 0.5, 0.5]

wh = 64

data_transform={
    
    "train":transforms.Compose([
            transforms.Resize((wh,wh)),
            transforms.RandomRotation(degrees=15),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean_nums, std_nums)
    
        ]),
    
    "test":transforms.Compose([
            
            transforms.Resize((wh, wh)),
            transforms.ToTensor(),
            transforms.Normalize(mean_nums, std_nums)
    
        ])
}


'''获取数据集路径'''

'''如果有一个组件是一个绝对路径，则在它之前的所有组件均会被舍弃'''

data_root=os.path.abspath(os.path.join(os.getcwd(), "."))

image_root=data_root+"/dataset/new_data/"


'''获取数据'''

train_dataset=torchvision.datasets.ImageFolder(root=image_root+"train",transform=data_transform["train"])

val_dataset=torchvision.datasets.ImageFolder(root=image_root+"test",transform=data_transform["test"])
    
'''设置参数'''

batch_size=32

epoch_total=1

class_num=3

learning_rate=0.001

'''装载数据'''

train_loader=torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)

val_loader=torch.utils.data.DataLoader(dataset=val_dataset,batch_size=batch_size,shuffle=False)

'''搭建神经网络'''

model=newsCNN.NewsNet()

summary.summary(model, input_size=(3, wh, wh),batch_size=batch_size,device="cpu")

'''损失函数'''

loss_function=nn.CrossEntropyLoss()

'''优化器'''

optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)
    
'''开始训练'''

step_total=len(train_loader)

for epoch in range(epoch_total):
    
    for step,(image,label) in enumerate(train_loader):
        
        pred=model(image)
        
        loss=loss_function(pred,label)
        
        optimizer.zero_grad()
        
        loss.backward()
        
        optimizer.step()
        
        if (step+1) % 100 == 0:
            
            print("Epoch:[{}/{}],Step:[{}/{}],Loss:{:.4f}".format(epoch, epoch_total,step+1,step_total,loss.item()))

save_file = os.path.join(os.path.dirname(__file__), '../model/news_classification_20220421.pkl')
torch.save(model, save_file)  
with torch.no_grad():
    
    correct=0
    
    total=0
    
    for image,label in val_loader:
        
        pred=model(image)
    
        predict=torch.max(pred,1)[1]
        
        correct += (predict == label).sum().item()
        
        total+=label.shape[0]
        
    print('Test Accuracy of the model on the  test images: {} %'.format(100 * correct / total))
